package com.xiaohan.turtlesoup.manager;

import cn.hutool.core.collection.CollUtil;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChoice;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * @program: TurtleSoup
 * @description: 客户端调用ai服务的方法封装【AI调用工具类】
 * @author: 韩小豪
 * @create: 2025-03-22 13:57
 **/

@Service
@Slf4j
public class AiManager {
	@Resource
	private ArkService arkService;

	/**
	 * 聊天时的ai服务调用（只允许传入系统预设和用户预设）
	 *
	 * @param sysPrompt  给ai的系统提示词【让ai知道自己要做什么，必须】
	 * @param userPrompt 用户提示词【用户向AI提的问题】
	 * @return ai思考后返回的结果
	 */
	public String doChat(String sysPrompt, String userPrompt) {
		/**
		 * 创建要发送给大模型的消息
		 */
		final List<ChatMessage> messages = new ArrayList<>();
		final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(sysPrompt).build();
		final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userPrompt).build();
		messages.add(systemMessage);
		messages.add(userMessage);

		//单次调用【给问题AI，ai组织好了答案后一次性返回】
		ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
				/**
				 * AI接入方式
				 * 1.直接指定大模型或接入点即可，更灵活，推荐开发AI类项目时使用
				 * 2.先在火山方舟中创建AI应用(机器人)，随后指定应用(机器人)id，更简单
				 * 3.将需要传递给ai的消息通过arkService对象发送
				 * 4.获取响应的结果并返回
				 */
				// 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
				.model("ep-20250321203745-jlbmp").messages(messages).build();

		List<ChatCompletionChoice> choices = arkService.createChatCompletion(chatCompletionRequest).getChoices();
		if (CollUtil.isEmpty(choices)) {
			throw new RuntimeException("AI服务调用异常，请重试...");
		}
		//由于项目需求一般基于一问一答的形式，而不需要列举，因此无需遍历获取所有的相应结果，只取第一个即可
				//forEach(choice -> System.out.println(choice.getMessage().getContent()));

		//取第一个回答并返回
		String aiResult = (String)choices.get(0).getMessage().getContent();//获取ai响应的消息中的具体答案内容，随后转换为字符串返回
		log.info(aiResult);
		return aiResult;
	}

	/**
	 * 更通用的ai调用方法封装，可以向AI传入任意条信息
	 * 目的：用以实现AI与用户交流间的上下文关联。【用户每次新的提问需要带上之前所有和AI的聊天记录，避免AI每次都是从0进行思考】
	 *
	 * @param chatMessageList 含之前所有的聊天记录
	 * @return
	 */
	public String doChat(List<ChatMessage> chatMessageList) {

		ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
				.model("ep-20250321203745-jlbmp").messages(chatMessageList).build();

		List<ChatCompletionChoice> choices = arkService.createChatCompletion(chatCompletionRequest).getChoices();
		if (CollUtil.isEmpty(choices)) {
			throw new RuntimeException("AI服务调用异常，请重试...");
		}

		String aiResult = (String)choices.get(0).getMessage().getContent();//获取ai响应的消息中的具体答案内容，随后转换为字符串返回
		log.info(aiResult);
		return aiResult;
	}

	/**
	 * 流式调用：传入一条消息
	 * 通过StringBuilder获取所有输出的结果
	 */
	public String doChatByStream(String sysPrompt, String userPrompt) {
		final List<ChatMessage> streamMessages = new ArrayList<>();
		final ChatMessage streamSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(sysPrompt).build();
		final ChatMessage streamUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userPrompt).build();
		streamMessages.add(streamSystemMessage);
		streamMessages.add(streamUserMessage);

		ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
				// 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
				.model("ep-20250321203745-jlbmp")
				.messages(streamMessages)
				.build();

		StringBuilder sb = new StringBuilder();
		arkService.streamChatCompletion(streamChatCompletionRequest)
				.doOnError(Throwable::printStackTrace)
				.blockingForEach(
						choice -> {
							if (choice.getChoices().size() > 0) {
								System.out.print(choice.getChoices().get(0).getMessage().getContent());
								sb.append(choice.getChoices().get(0).getMessage().getContent());
							}
						}
				);
		return sb.toString();  //将所有流式输出结果转换为字符串返回
	}

	/**
	 * 流式调用：传入消息列表
	 * @param chatMessageList 含历史对话记录的消息列表
	 */
	public String doChatByStream(List<ChatMessage> chatMessageList) {

		ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
				// 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
				.model("ep-20250321203745-jlbmp")
				.messages(chatMessageList)
				.build();

		StringBuilder sb = new StringBuilder();
		arkService.streamChatCompletion(streamChatCompletionRequest)
				.doOnError(Throwable::printStackTrace)
				.blockingForEach(
						choice -> {
							if (choice.getChoices().size() > 0) {
								System.out.print(choice.getChoices().get(0).getMessage().getContent());
								sb.append(choice.getChoices().get(0).getMessage().getContent());
							}
						}
				);
		return sb.toString();  //将所有流式输出结果转换为字符串返回
	}
}
